import logging as log
import re
import google.generativeai as genai
from google.ai import generativelanguage as glm
from google.generativeai.generative_models import GenerativeModel
from prompt_toolkit import prompt
log.basicConfig(level=log.DEBUG)

def get_api_key(key:str="GOOGLE_API_KEY"):
    try:
        import os
        GOOGLE_API_KEY = os.environ[key]
    except ImportError:
        log.error("从环境变量中获取GOOGLE_API_KEY失败")
    log.debug(f"GOOGLE_API_KEY：{GOOGLE_API_KEY}")
    return GOOGLE_API_KEY
    
MODEL_INSTRUCTIONS="""\
#角色定义
    你是一个高级python专家，你将根据用户的需求生成python程序。

#功能定义
    将用户的需求转换成python程序,python程序要求如下：
    ## 代码格式：
        1.使用中文注释，每行最多80个字符。
        2.导入包顺序依次是：系统库、第三方库、自定义库；然后按字母排序。
        3.函数名、变量名、类名按照python编程规范定义，做到简短，含义精准。
        4.禁止生成文档注释！禁止生成文档注释！禁止生成文档注释！
    ## 代码实现：
        1.请分析程序是否是IO密集型任务，如果是使用asyncio库，提高程序性能。
        2.如果需要网络请求请优先使用aiohttp库。
        3.如果需要文件操作请优先使用aiofile库。
        4.如果需要持久化数据请使用aiomysql、asyncpg 库，如果用户没有指定数据库默认使用asyncpg库。
        5.如果需要提取页面内容请使用beautifulsoup库
""" 
def get_gemini_model(google_api_key:str,
                     model_name:str="gemini-1.5-pro-latest"):
    genai.configure(api_key=google_api_key,transport="rest")
    model = genai.GenerativeModel(model_name=model_name,
                        system_instruction=MODEL_INSTRUCTIONS)
    log.debug(f"获取当前模型：{model}")
    return model

def user_chat(prompt:str,
                    model:GenerativeModel,
                    chat_history:bool=True):
    #创建聊天对象
    chat = model.start_chat()
    if chat_history:
        for content in chat.history:
            role = content.role
            chat_content = [type(part).to_dict(part) for part in content.parts]
            log.info(f"聊天记录：{role}=>{chat_content}")
    response = chat.send_message(content=prompt)
    log.info(f"获取响应内容：{response}")
    return response


def extract_python_code_blocks(markdown_text:str):
    """从markdown文档中提取python代码块"""
    # 正则表达式模式，用于匹配Markdown代码块
    python_code_block_pattern = r'```python\n(.*?)\n```'
    matches = re.findall(python_code_block_pattern, 
                        markdown_text, 
                        re.DOTALL | re.MULTILINE)
    python_code_blocks=""
    for match in matches:
        python_code_blocks += match.strip()
    if len(python_code_blocks)==0:
        raise ValueError("markdown文档中没有python脚本。")
    log.debug(f"从markdown文档中提取的python代码块：{python_code_blocks}")
    return python_code_blocks

def exec_python_script(python_script:str):
    """执行python 脚本"""
    exec(python_script, globals())


exec_python_script_tool = glm.Tool(
    function_declarations=[
        glm.FunctionDeclaration(
            name="exec_python_script",
            description="执行python脚本",
            parameters=glm.Schema(
                type=glm.Type.STRING,
                required=["python_script"]
            )
        )

    ]
)
print(exec_python_script_tool.function_declarations)
api_key = get_api_key()
model = get_gemini_model(api_key)

# prompt="请帮助我用python实现输出 'hello world' 这个字符串在txt文件上。"
prompt="请帮我爬取url：https://www.qq.com/ 地址，然后提取页面文本储在txt文件。\
    注意：编码使用utf-8；提取文本可以使用beautifulsoup"
response = user_chat(prompt=prompt,
          model=model)
markdown_text=response.candidates[0].content.parts[0].text
python_script = extract_python_code_blocks(markdown_text)
#执行生成的python脚本。
exec_python_script(python_script)
# chat = model.start_chat(enable_automatic_function_calling=True)
# resp = chat.send_message(python_script,tools=[exec_python_script_tool])
# log.info(f"执行结果：{resp.candidates[0].content.parts[0].text}")